from zipline.api import order_target_percent, symbol, record
import numpy as np

def initialize(context):
    """
    初始化动量策略参数
    """
    # 策略参数
    context.long_period = 252  # 12个月收益率计算周期
    context.short_period = 21  # 1个月收益率计算周期
    context.momentum_threshold = 0.05  # 动量强度触发阈值
    context.max_sector_exposure = 0.25  # 单一行业持仓上限
    context.asset = symbol('600519')  # 使用A股标的：贵州茅台
    
    # 风险控制参数
    context.volatility_threshold = 0.4  # 波动率过滤阈值

def handle_data(context, data):
    """
    处理每日数据，实现动量策略核心逻辑
    """
    # 获取历史价格数据
    prices = data.history(
        context.asset, 
        'price', 
        bar_count=max(context.long_period, context.short_period) + 1, 
        frequency='1d'
    )
    
    # 计算波动率（风险控制）
    volatility = np.std(prices[-30:])  # 30日波动率
    if volatility > context.volatility_threshold:
        # 波动率过高，暂停交易
        return
    
    # 计算长期和短期收益率
    long_return = (prices[-1] / prices[-context.long_period-1]) - 1
    short_return = (prices[-1] / prices[-context.short_period-1]) - 1
    
    # 当前持仓
    position = context.portfolio.positions[context.asset].amount
    
    # 交易信号逻辑
    if long_return > short_return + context.momentum_threshold:
        # 动量增强，增加仓位
        if position == 0:  # 无持仓时开仓
            order_target_percent(context.asset, 0.05)
    elif long_return < short_return - context.momentum_threshold:
        # 动量减弱，减少仓位
        if position != 0:  # 有持仓时减仓
            order_target_percent(context.asset, 0)
    
    # 行业分散控制（简化示例）
    # 实际实现需要获取行业信息并计算行业敞口
    # 这里仅作为占位符
    sector_exposure = 0.0  # 实际实现中计算行业敞口
    if sector_exposure > context.max_sector_exposure:
        # 超过行业上限，减仓
        order_target_percent(context.asset, 0)
    
    # 记录数据（可选）
    record(
        price=prices[-1],
        long_return=long_return,
        short_return=short_return,
        volatility=volatility
    )